import numpy as np
import pandas as pd


def ind_cnt(yi, xi, ti, vi, di):
    xi[xi < 0] = 0
    xi[xi >= 330] = 329
    yi[yi < 0] = 0
    yi[yi >= 110] = 109
    df_t = pd.DataFrame([xi, yi, ti], index=['x', 'y', 't']).T.groupby(['x', 'y'])['t'].agg({'c_cnt': 'count'}).reset_index()
    df_v = pd.DataFrame([xi, yi, vi], index=['x', 'y', 'v']).T.groupby(['x', 'y'])['v'].agg({'v_mean': 'mean'}).reset_index()
    df_d = pd.DataFrame([xi, yi, di], index=['x', 'y', 'd']).T.groupby(['x', 'y'])['d'].agg({'d_mean': 'mean'}).reset_index()
    return df_t, df_v, df_d


train = pd.read_hdf('./input/train.h5')
arr = []
ships = train.loc[:, 'ship'].unique()
for i in range(len(ships)):
    I = np.zeros((110, 330, 3))
    ship = ships[i]
    data0 = train.loc[(train['ship'] == ship) & (train['ship'] == ship), :]
    lon, lat = data0['y'] / 50000, data0['x'] / 200000
    xi = ((lon - 90) / 0.1).astype(np.int)
    yi = ((lat - 25) / 0.1).astype(np.int)
    ti, vi, di = data0['time'], data0['v'], data0['d']
    W, df_v, df_d = ind_cnt(yi, xi, ti, vi, di)
    print(i, data0['type'][0], ': 停驻位置个数', len(W), ',', 'max_cnt:', max(W.iloc[:, 2]),
          'max_v:', max(df_v.iloc[:, 2]), 'max_d:', max(df_v.iloc[:, 2]))
    xi1, yi1 = W['x'], W['y']
    for j, a in enumerate([W, df_v, df_d]):
        I1 = list(a.iloc[:, 2])
        I2 = I1/np.max(I1)*255
        I3 = I2.astype(np.uint8)
        I[yi1, xi1, j] = I3
    arr += [I.reshape(-1)]

Arr = np.array(arr)


res = pd.DataFrame(Arr)
type_map = dict(zip(train['type'].unique(), np.arange(3)))
res['label'] = train.drop_duplicates(['ship', 'type'])['type'].reset_index()['type'].map(type_map)
res['ship'] = train.drop_duplicates(['ship', 'type']).reset_index()['ship']
res.to_hdf('./data/im1x3_tvd.h5', 'df', mode='w')
